from fastapi.openapi.models import Response
from langchain import PromptTemplate, LLMChain
import uvicorn
from fastapi import FastAPI
from langchain.chains.base import Chain
from langchain_community.chat_models import ChatTongyi
from langchain_core.prompts import ChatPromptTemplate
from starlette.responses import StreamingResponse

app = FastAPI()
import os

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"


# 定义一个生成器函数，用于流式输出结果
def stream_output(query):
    # 初始化OpenAI LLM
    llm = ChatTongyi(stream=True)
    # 创建一个PromptTemplate，用于生成提示
    prompt = ChatPromptTemplate.from_messages(
        [("system", "你是一个专业的AI助手。"), ("human", "{query}")]
    )
    # 创建LLM链，用于处理用户的查询并流式输出结果
    llm_chain = LLMChain(llm=llm, prompt=prompt)
    # llm_chain = prompt | llm
    reponse = llm_chain.stream({"query": query})
    for res in reponse:
        yield res.content


# # 假设这是一个Web服务器端的函数，用于处理HTTP请求并流式响应
def web_stream_response(query):
    # 这里可以是你的Web框架特定的代码，比如Flask的response流式响应
    for chunk in stream_output(query):
        # 输出每一块数据到前端，例如使用Flask的response对象
        print(chunk, end="", flush=True)
        yield chunk


# 使用方式
# 在Web应用中，你可以将web_stream_response函数与HTTP请求结合起来使用
# 例如在Flask中，可以这样使用:
@app.get("/")
def stream_output_view():
    response = web_stream_response("python和Java的区别是什么，请详细解释")
    # StreamingResponse返回一个流式响应对象，模拟sse协议
    return StreamingResponse(response, media_type="text/event-stream")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=9001)
